import numpy as np
from PIL import Image

# 尝试导入 rembg，若缺失则在调用时抛错提示用户安装依赖
try:
    from rembg import remove  # type: ignore
except ImportError as e:  # noqa: PERF203,F401
    remove = None  # type: ignore


def u2net_alpha(rgb_img: np.ndarray) -> np.ndarray:
    """使用 rembg (U²-Net) 生成 alpha 通道

    参数
    -----
    rgb_img : np.ndarray
        形状 (H, W, 3) 的 RGB 图像数据，dtype=uint8。

    返回
    -----
    np.ndarray
        形状 (H, W) 的 alpha 蒙版 (uint8)，取值 0-255。
    """
    if remove is None:  # pragma: no cover – 环境缺依赖
        raise RuntimeError("未安装 rembg，请先运行:  pip install rembg onnxruntime")

    # Pillow 对象
    pil_in = Image.fromarray(rgb_img)

    # rembg.remove 可直接接受 PIL.Image，返回 PIL.Image (RGBA)
    pil_out = remove(pil_in)

    out_np = np.asarray(pil_out)

    # 若输出不包含透明通道，则直接返回全不透明
    if out_np.ndim == 2:  # 灰度
        alpha = (out_np > 0).astype(np.uint8) * 255
    elif out_np.shape[2] == 4:
        alpha = out_np[:, :, 3]
    else:
        # 退路：把所有非白像素视作前景
        alpha = 255 - ((out_np[:, :, :3] > 240).all(axis=2).astype(np.uint8) * 255)

    return alpha.astype(np.uint8) 